#include "../common/typedef.h"
#include <cassert>
#include <cmath>

namespace kuai {

	void
		add_e_n12_6(
			const RealNumber coords[], const RealNumber k[],
			const ItemIndex items[], int n,
			RealNumber energies[]
		)
	{
		assert (n % 3 == 0);
		const ItemIndex* p0 = items;
		const ItemIndex* p1 = items+n;
		while (p0 < p1) {
			AtomIndex a1 = *(p0++);
			AtomIndex a2 = *(p0++);
			ItemIndex ki = *(p0++);
			RealNumber x = coords[a1*3+0]-coords[a2*3+0];
			RealNumber y = coords[a1*3+1]-coords[a2*3+1];
			RealNumber z = coords[a1*3+2]-coords[a2*3+2];

			RealNumber krr2 = k[ki] / (x*x+y*y+z*z);
			RealNumber krr6 = krr2 * krr2 * krr2;

			energies[0] += k[ki+1] * krr6 * (krr6 - 2);
		}
	}

	void
		add_f_n12_6(
			const RealNumber coords[], const RealNumber k[],
			const ItemIndex items[], int n,
			RealNumber energies[], RealNumber forces[]
		)
	{
		assert (n % 3 == 0);
		const AtomIndex* p0 = items;
		const AtomIndex* p1 = items+n;
		while (p0 < p1) {
			AtomIndex a1 = *(p0++);
			AtomIndex a2 = *(p0++);
			ItemIndex ki = *(p0++);
			RealNumber x = coords[a1*3+0]-coords[a2*3+0];
			RealNumber y = coords[a1*3+1]-coords[a2*3+1];
			RealNumber z = coords[a1*3+2]-coords[a2*3+2];

			RealNumber rr2 = 1 / (x*x+y*y+z*z);
			RealNumber krr2 = k[ki] * rr2;
			RealNumber krr6 = krr2 * krr2 * krr2;

			energies[0] += k[ki+1] * krr6 * (krr6 - 2);

			RealNumber factor = 12 * k[ki+1] * krr6 * (krr6-1) * rr2;
			x *= factor, y *= factor, z *= factor;
			forces[a1*3+0] += x, forces[a1*3+1] += y, forces[a1*3+2] += z;
			forces[a2*3+0] -= x, forces[a2*3+1] -= y, forces[a2*3+2] -= z;
		}
	}

	void
		add_e_coulomb(
			const RealNumber coords[], const RealNumber q[], const RealNumber COULOMB_FACTOR,
			const AtomIndex items[], int n,
			RealNumber energies[]
		)
	{
		assert (n % 3 == 0);
		const AtomIndex* p0 = items;
		const AtomIndex* p1 = items+n;
		while (p0 < p1) {
			AtomIndex a1 = *(p0++);
			AtomIndex a2 = *(p0++);
			p0++;
			RealNumber x = coords[a1*3+0]-coords[a2*3+0];
			RealNumber y = coords[a1*3+1]-coords[a2*3+1];
			RealNumber z = coords[a1*3+2]-coords[a2*3+2];

			RealNumber rr = 1 / sqrt(x*x+y*y+z*z);
			energies[0] += COULOMB_FACTOR * q[a1] * q[a2] * rr;
		}
	}

	void
		add_f_coulomb(
			const RealNumber coords[], const RealNumber q[], const RealNumber COULOMB_FACTOR,
			const AtomIndex items[], int n,
			RealNumber energies[], RealNumber forces[]
		)
	{
		assert (n % 3 == 0);
		const AtomIndex* p0 = items;
		const AtomIndex* p1 = items+n;
		while (p0 < p1) {
			AtomIndex a1 = *(p0++);
			AtomIndex a2 = *(p0++);
			p0++;
			RealNumber x = coords[a1*3+0]-coords[a2*3+0];
			RealNumber y = coords[a1*3+1]-coords[a2*3+1];
			RealNumber z = coords[a1*3+2]-coords[a2*3+2];

			RealNumber rr2 = 1 / (x*x+y*y+z*z);
			RealNumber rr = sqrt(rr2);

			RealNumber factor = COULOMB_FACTOR * q[a1] * q[a2] * rr;

			energies[0] +=  factor;
			factor *= rr2;

			x *= factor, y *= factor, z *= factor;
			forces[a1*3+0] += x, forces[a1*3+1] += y, forces[a1*3+2] += z;
			forces[a2*3+0] -= x, forces[a2*3+1] -= y, forces[a2*3+2] -= z;
		}
	}
}
